import cv2
import numpy as np
from argparse import ArgumentParser


def add_gaussian_noise(img, mu=0, sigma=1):
    result = np.zeros(img.shape, dtype=np.uint8)
    x, y = img.shape
    for i in range(x):
        for j in range(y):
            result[i][j] = img[i][j] + int(np.random.normal(loc=mu, scale=sigma))
    return result


def main(option):
    np.random.seed(option.random_seed)
    img = cv2.imread('lena_gray.bmp', cv2.IMREAD_GRAYSCALE)
    img = add_gaussian_noise(img, mu=option.mu, sigma=option.sigma)
    cv2.imwrite(option.output_file % (option.mu, option.sigma), img)
    # cv2.namedWindow('Image')
    # cv2.imshow('Image', img)
    # cv2.waitKey()


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--random_seed', type=int, default=19950125)
    parser.add_argument('--mu', type=float, default=0)
    parser.add_argument('--sigma', type=float, default=1)
    parser.add_argument('--output_file', type=str, default='output/noise/gaussian_mu=%.2f_sigma=%.2f.bmp')
    option = parser.parse_args()
    main(option)
